
#include <stdio.h>
#include <stdlib.h>
#include "../iface/ZlibInterface.h"
#include "GZCompress.h"

/** Construct the class with no bucket brigades. This means that this
 * is to be used for block and not pipe compression. */
GZCompress::GZCompress()
{
   input = NULL;
   output = NULL;
   strm = NULL;
}

/** Construct the class with bucket brigades. This means that this can
 * be used for both block and pipe compression.
 * 
 * @param iinput The bucket brigade for input.
 * 
 * @param ioutput The bucket brigade for output. */
GZCompress::GZCompress(APRBucketBrigade &iinput, APRBucketBrigade &ioutput)
{
   input = &iinput;
   output = &ioutput;
   strm = NULL;
}

/** Destructor for the zlib compression class. */
GZCompress::~GZCompress()
{
}

/** Compress data block and return compressed data. The compressed
 * data returned must be freed by the caller.
 * 
 * @param inputblock A block of plaintext data to compress
 * (compressing encrypted data doesn't really work).
 * 
 * @param inputsize The length of the plaintext data block to
 * compress.
 * 
 * @param[out] outputsize The length of the compressed block of data
 * returned. The data returned must be freed by the caller. */
char *GZCompress::processDataBlock(const char *inputblock, apr_size_t inputsize, apr_size_t *outputsize)
{
   apr_size_t max_output_size = compressBound(inputsize);
   uLongf dest_length = max_output_size;
   char *result = (char *)malloc(sizeof(char)*max_output_size);
   int zstatus = compress((Bytef *)result, &dest_length, (const Bytef *)inputblock, inputsize);
   CheckZlibError(zstatus);
   if (dest_length > 0)
   {
      if (outputsize)
         *outputsize = dest_length;
      return result;
   }
   free(result);
   return NULL;
}

/** Begin the process of compressing a pipe of data.
 * 
 * @return Zero on success, non-zero on failure. */
int GZCompress::processPipeBegin()
{
   if (strm)
   {
      errprint("already in use");
      return 1;
   }
   int zstatus;
   strm = (z_stream *)malloc(sizeof(z_stream));

   strm->zalloc = Z_NULL;
   strm->zfree = Z_NULL;
   strm->opaque = Z_NULL;
   
   zstatus = deflateInit(strm, Z_DEFAULT_COMPRESSION);
   CheckZlibError(zstatus);
   if (zstatus != Z_OK) return zstatus;
   return 0;
}

/** Process any outstanding data on the input pipe, and place any
 * processed compressed data into the output pipe.
 * 
 * @return Zero on success, non-zero on failure. */
int GZCompress::processPipeStep()
{
   unsigned char *itemp_block = NULL;
   unsigned char *otemp_block = NULL;
   apr_size_t input_block = input->realLength();
   apr_size_t readbytes = 0;
   apr_size_t writebytes = 0;
   apr_size_t wrotebytes = 0;
   int retval = 1;
   int zstatus = Z_OK;
   while (input_block != 0)
   {
      itemp_block = (unsigned char *)malloc(sizeof(unsigned char)*input_block);
      otemp_block = (unsigned char *)malloc(sizeof(unsigned char)*input_block);
      readbytes = input->readFromStream((char *)itemp_block, input_block);
      if (readbytes != input_block)
      {
         errprint("failed to read data from plain bucket brigade (read %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), buffer unfinished", readbytes, input_block);
         goto stopsteps;
      }
   
      strm->avail_in = input_block;
      strm->next_in = itemp_block;
      strm->avail_out = input_block;
      strm->next_out = otemp_block;
      zstatus = deflate(strm, Z_NO_FLUSH);
      CheckZlibError(zstatus);
      writebytes = input_block - (strm->avail_out);
      wrotebytes = output->writeToStream((const char *)otemp_block, writebytes);
      if (wrotebytes != writebytes)
      {
         errprint("failed to write data to compressed bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", wrotebytes, writebytes);
         goto stopsteps;
      }
   
      input->removeFromStream(input_block - (strm->avail_in));
      if (itemp_block) free(itemp_block);
      itemp_block = NULL;
      if (otemp_block) free(otemp_block);
      otemp_block = NULL;
      input_block = input->realLength();
   }
   retval = 0;
stopsteps:
   if (itemp_block) free(itemp_block);
   itemp_block = NULL;
   if (otemp_block) free(otemp_block);
   otemp_block = NULL;
   return 0;
}

/** Process the pipe step and record the movement of data to determine
 * when the process is complete.
 * 
 * @param[out] movementin Records data input in bytes.
 * 
 * @param[out] movementout Records data output in bytes. */
int GZCompress::processPipeStep(int *movementin, int *movementout)
{
   apr_size_t before_input = input->realLength();
   apr_size_t before_output = output->realLength();
   int retval = processPipeStep();
   apr_size_t after_input = input->realLength();
   apr_size_t after_output = output->realLength();
   if (movementin)
      *movementin = before_input - after_input;
   if (movementout)
      *movementout = after_output - before_output;
   return retval;
}

/** Process any outstanding plaintext data in the pipe and finalise
 * the compression process.
 * 
 * @return Zero on success, non-zero on failure. */
int GZCompress::processPipeEnd()
{
   processPipeStep();
   const int blocksize = 0x10000;
   unsigned char *otemp_block = (unsigned char *)malloc(sizeof(unsigned char)*blocksize);
   do
   {
      strm->avail_in = 0;
      strm->next_in = NULL;
      strm->avail_out = blocksize;
      strm->next_out = otemp_block;
      int zstatus = deflate(strm, Z_FINISH);
      CheckZlibError(zstatus);
      apr_size_t writebytes = blocksize - (strm->avail_out);
      apr_size_t wrotebytes = output->writeToStream((const char *)otemp_block, writebytes);
      if (wrotebytes != writebytes)
      {
         errprint("failed to write data to compressed bucket brigade (wrote %" APR_SIZE_T_FMT ", had %" APR_SIZE_T_FMT "), likely corrupt", wrotebytes, writebytes);
         deflateEnd(strm);
         free(strm);
         strm = NULL;
         free(otemp_block);
         return 1;
      }
   } while (strm->avail_out == 0);
   deflateEnd(strm);
   free(strm);
   strm = NULL;
   free(otemp_block);
}
